#include "ML_XMLParsers.h"
#include "ImageFragment.h"
#include "ImageLoader.h"
#include "ML_JointBoosting.h"
#include "ML_JointBoostingParams.h"
#include "ML_DecisionStump.h"
#include "ML_DecisionStumpParams.h"
#include "../Scene/ZDataManager.h"

GLOBAL_NAMESPACE_BEGIN
NAMESPACE_MACHINE_LEARNING_BEGIN

const char* MLXmlParsers::ElementNames[] = {
    "fragmentParams",
    "fragment",
    "JointBoosting",
    "decisionStump",
    "imageLoader"
};

#undef TINYXML2_CHECK_VALID_AND_OUTPUT
#define TINYXML2_CHECK_VALID_AND_OUTPUT(ret) if((ret)!=tinyxml2::XML_NO_ERROR) return NULL;

tinyxml2::XMLElement* MLXmlParsers::parseToXMLElement(tinyxml2::XMLDocument* doc, const ImageFragmentParams* param)
{
    tinyxml2::XMLElement* elem = doc->NewElement(ElementNames[ML_XML_ENT_ImageFragmentParams]);
    tinyxml2::XMLElement* window = doc->NewElement("window");
    window->SetAttribute("x", param->u);
    window->SetAttribute("y", param->v);
    window->SetAttribute("w", param->w);
    window->SetAttribute("h", param->h);
    tinyxml2::XMLElement* srcChild = doc->NewElement("folder");
    srcChild->SetText(param->fragmentPath.c_str());
    window->LinkEndChild(srcChild);
    elem->LinkEndChild(window);
    return elem;
}

tinyxml2::XMLElement* MLXmlParsers::parseToXMLElement(tinyxml2::XMLDocument* doc, const ImageFragment* fragment)
{
    tinyxml2::XMLElement* elem = doc->NewElement(ElementNames[ML_XML_ENT_ImageFragment]);
    tinyxml2::XMLElement* imaPara  = doc->NewElement("image");
    imaPara->SetAttribute("src",         fragment->getImage()->getFullPath().c_str());
    imaPara->SetAttribute("id",          fragment->getImage()->getId());
    imaPara->SetAttribute("templateSrc", fragment->getTemplateImage()->getFullPath().c_str());
    imaPara->SetAttribute("tid",         fragment->getTemplateImage()->getId());
    imaPara->SetAttribute("maskSrc",     fragment->getMaskImage()->getFullPath().c_str());
    imaPara->SetAttribute("mid",         fragment->getMaskImage()->getId());
    imaPara->SetAttribute("order",       fragment->getOrder());
    tinyxml2::XMLElement* fragPara = parseToXMLElement(doc, fragment->getFragmentParams());
    elem->LinkEndChild(imaPara);
    elem->LinkEndChild(fragPara);
    return elem;
}

tinyxml2::XMLElement* MLXmlParsers::parseToXMLElement(tinyxml2::XMLDocument* doc, const MLJointBoosting* jbAlgo)
{
    if (!jbAlgo) return NULL;
    tinyxml2::XMLElement* elem = doc->NewElement(ElementNames[ML_XML_ENT_MLJointBoosting]);
    elem->SetAttribute("roundM", jbAlgo->getParams()->roundM);
    elem->SetAttribute("classSize", jbAlgo->getParams()->classSize);
    for (int i = 0; i < jbAlgo->getWeakLearnerCount(); i++)
    {
        MLDecisionStump* stump = dynamic_cast<MLDecisionStump*>(jbAlgo->getWeakLearnerOf(i));
        if (stump)
        {
            tinyxml2::XMLElement* oneStump = parseToXMLElement(doc, stump);
            elem->LinkEndChild(oneStump);
        }
    }
    return elem;
}

tinyxml2::XMLElement* MLXmlParsers::parseToXMLElement(tinyxml2::XMLDocument* doc, const MLDecisionStump* stump)
{
    if (!stump) return NULL;
    tinyxml2::XMLElement* elem = doc->NewElement(ElementNames[ML_XML_ENT_MLDecisionStump]);
    const MLDecisionStumpParams* param = stump->getParameters();
    tinyxml2::XMLElement* subSet = doc->NewElement("subSet");
    tinyxml2::XMLElement* ks = doc->NewElement("k");
    std::stringstream ss;
    for (int i = 0; i < param->subClass.size(); i++)
    {
        ss.str("");
        ss << param->subClass.at(i);
        tinyxml2::XMLElement* oneSet = doc->NewElement("p");
        oneSet->SetText(ss.str().c_str());
        subSet->LinkEndChild(oneSet);
    }
    for (int i=0; i<param->k.size(); i++)
    {    
        ss.str("");
        ss << param->k.at(i);
        tinyxml2::XMLElement* oneSet = doc->NewElement("p");
        oneSet->SetText(ss.str().c_str());
        ks->LinkEndChild(oneSet);
    }
    elem->LinkEndChild(subSet);
    elem->LinkEndChild(ks);
   //elem->SetAttribute("subSet",    ss.str().c_str());
    elem->SetAttribute("dim",       param->dimF);
    elem->SetAttribute("greater",   param->bGreater);
    elem->SetAttribute("threshold", param->threshold);
    elem->SetAttribute("a",         param->a);
    elem->SetAttribute("b",         param->b);
    //elem->SetAttribute("k",         ssk.str().c_str());
    return elem;
}

void* MLXmlParsers::parseFromXMLElement(tinyxml2::XMLElement* elem)
{
    std::string name(elem->Name());
    if (name.compare(ElementNames[ML_XML_ENT_ImageFragmentParams])==0)
    {
        ImageFragmentParams* ret = new ImageFragmentParams;
        if (parseFromXMLElement(elem, ret)!=NULL)
            return (void*)ret;
        else
        {
            delete ret;
            return NULL;
        }
    }
    else if (name.compare(ElementNames[ML_XML_ENT_ImageFragment])==0)
    {
        ImageFragment* ret = new ImageFragment("");
        if (parseFromXMLElement(elem, ret)!=NULL)
            return (void*)ret;
        else
        {
            delete ret;
            return NULL;
        }
    }
    else if (name.compare(ElementNames[ML_XML_ENT_MLJointBoosting])==0)
    {
        MLJointBoosting* ret = new MLJointBoosting;
        if (parseFromXMLElement(elem, ret)!=NULL)
            return (void*)ret;
        else
        {
            delete ret;
            return NULL;
        }
    }
    else if (name.compare(ElementNames[ML_XML_ENT_MLDecisionStump])==0)
    {
        MLDecisionStump* ret = new MLDecisionStump;
        if (parseFromXMLElement(elem, ret)!=NULL)
            return (void*)ret;
        else
        {
            delete ret;
            return NULL;
        }
    }
    else if (name.compare(ElementNames[ML_XML_ENT_ImageLoader])==0)
    {
        ImageLoader* ret = new ImageLoader(DataManager::getSingletonPtr());
        if (parseFromXMLElement(elem, ret)!=NULL)
            return (void*)ret;
        else
        {
            delete ret;
            return NULL;
        }

    }
    else
    {
        return NULL;
    }
}

ImageFragmentParams* MLXmlParsers::parseFromXMLElement(tinyxml2::XMLElement* elem, ImageFragmentParams* param)
{
    if (!elem || strcmp(ElementNames[ML_XML_ENT_ImageFragmentParams], elem->Value())!=0)
        return NULL;
    
    tinyxml2::XMLElement* windowElem = elem->FirstChildElement("window");
    if (!windowElem) return NULL;
    TINYXML2_CHECK_VALID_AND_OUTPUT(windowElem->QueryIntAttribute("x", &param->u));
    TINYXML2_CHECK_VALID_AND_OUTPUT(windowElem->QueryIntAttribute("y", &param->v));
    TINYXML2_CHECK_VALID_AND_OUTPUT(windowElem->QueryIntAttribute("w", &param->w));
    TINYXML2_CHECK_VALID_AND_OUTPUT(windowElem->QueryIntAttribute("h", &param->h));
    tinyxml2::XMLElement *srcChild = windowElem->FirstChildElement("folder");
    if (!srcChild) return NULL;
    param->fragmentPath = std::string(srcChild->Value());
    return param;
}

ImageFragment*       MLXmlParsers::parseFromXMLElement(tinyxml2::XMLElement* elem, ImageFragment* imageFragment)
{
    if (!elem || strcmp(ElementNames[ML_XML_ENT_ImageFragment], elem->Value())!=0)
        return NULL;

    ImageFragmentParamsIO io;
    tinyxml2::XMLElement* imageElem = elem->FirstChildElement("image");
    io.imaFullPath      = std::string(imageElem->Attribute("src"));
    io.templateFullPath = std::string(imageElem->Attribute("templateSrc"));
    io.maskFullPath     = std::string(imageElem->Attribute("maskSrc"));
    TINYXML2_CHECK_VALID_AND_OUTPUT(imageElem->QueryIntAttribute("id", &io.id));
    TINYXML2_CHECK_VALID_AND_OUTPUT(imageElem->QueryIntAttribute("tid", &io.templateId));
    TINYXML2_CHECK_VALID_AND_OUTPUT(imageElem->QueryIntAttribute("mid", &io.maskId));
    TINYXML2_CHECK_VALID_AND_OUTPUT(imageElem->QueryIntAttribute("order", &io.order));
    ImageFragmentParams* para = parseFromXMLElement(
        elem->FirstChildElement(ElementNames[ML_XML_ENT_ImageFragmentParams]), &io.fragmentPara);
    imageFragment->init(io);
    return imageFragment;
}

MLJointBoosting*     MLXmlParsers::parseFromXMLElement(tinyxml2::XMLElement* elem, MLJointBoosting* jbAlgo)
{
    if (!elem || strcmp(ElementNames[ML_XML_ENT_MLJointBoosting], elem->Value())!=0)
        return NULL;
    tinyxml2::XMLElement* child = elem->FirstChildElement(ElementNames[ML_XML_ENT_MLDecisionStump]);
    while (child)
    {
        MLDecisionStump *stump = new MLDecisionStump;
        if (parseFromXMLElement(child, stump)!=NULL)
        {
            jbAlgo->addWeakLearner(stump);
        }
        else
        {
            delete stump;
        }
        child = child->NextSiblingElement(ElementNames[ML_XML_ENT_MLDecisionStump]);
    }
    return jbAlgo;
}

MLDecisionStump*     MLXmlParsers::parseFromXMLElement(tinyxml2::XMLElement* elem, MLDecisionStump* stump)
{
    MLDecisionStumpParams params;
    std::string subSetStr = std::string(elem->Attribute("subSet"));;
    //std::stringstream ss;
    //ss << subSetStr;
    //int oneset = -1;
    //while(ss >> oneset)
    //{
    //    params.subClass.push_back(oneset);
    //}
    TINYXML2_CHECK_VALID_AND_OUTPUT(elem->QueryIntAttribute("dim", &params.dimF));
    TINYXML2_CHECK_VALID_AND_OUTPUT(elem->QueryBoolAttribute("greater", &params.bGreater));
    TINYXML2_CHECK_VALID_AND_OUTPUT(elem->QueryDoubleAttribute("threshold", &params.threshold));
    TINYXML2_CHECK_VALID_AND_OUTPUT(elem->QueryDoubleAttribute("a", &params.a));
    TINYXML2_CHECK_VALID_AND_OUTPUT(elem->QueryDoubleAttribute("b", &params.b));
    tinyxml2::XMLElement* subset = elem->FirstChildElement("subSet");
    if (!subset)
        return NULL;
    tinyxml2::XMLElement* oneset = subset->FirstChildElement("p");
    int s = -1;
    while (oneset)
    {
        tinyxml2::XMLUtil::ToInt(oneset->GetText(), &s);
        params.subClass.push_back(s);
        oneset = oneset->NextSiblingElement("p");
    }
    tinyxml2::XMLElement* ks = elem->FirstChildElement("k");
    if (!ks) return NULL;
    tinyxml2::XMLElement* oneK = ks->FirstChildElement("p");
    double k = 0;
    while (oneset)
    {
        tinyxml2::XMLUtil::ToDouble(oneset->GetText(), &k);
        params.k.push_back(k);
        oneset = oneset->NextSiblingElement("p");
    }
    /*std::string kStr = std::string(elem->Attribute("k"));
    std::stringstream ss1;
    ss1 << kStr;
    double oneK;
    while(ss1 >> oneK)
    {
        params.k.push_back(oneK);
    }*/
    if (params.subClass.empty() || params.k.empty())
        return NULL;
    stump->buildStump(params);
    return stump;
}

tinyxml2::XMLElement* MLXmlParsers::parseToXMLElement(tinyxml2::XMLDocument* doc, const ImageLoader* imageLoader)
{
    if (!imageLoader || !doc) return NULL;
    tinyxml2::XMLElement* elem = doc->NewElement(ElementNames[ML_XML_ENT_ImageLoader]);
    int nfrag = imageLoader->getFragmentCount();
    for (int i = 0; i < nfrag; i++)
    {
        const ImageFragment* oneFragment = imageLoader->getFragment(i);
        tinyxml2::XMLElement* child = parseToXMLElement(doc, oneFragment);
        elem->LinkEndChild(child);
    }
    return elem;
}

ImageLoader* MLXmlParsers::parseFromXMLElement(tinyxml2::XMLElement* elem, ImageLoader* imageLoader)
{
    if (!imageLoader || !elem) return NULL;
    tinyxml2::XMLElement* child = elem->FirstChildElement(ElementNames[ML_XML_ENT_ImageFragment]);
    while (child)
    {
        ImageFragment frag;
        ImageFragment* ret = parseFromXMLElement(child, &frag);
        if (ret)
            imageLoader->addFragment(frag);
        child = child->NextSiblingElement(ElementNames[ML_XML_ENT_ImageFragment]);
    }
    return imageLoader;
}

NAMESPACE_MACHINE_LEARNING_END
GLOBAL_NAMESPACE_END
